# -*- coding:utf-8 -*-
"""
# @file name    : my_work.py
# @author       : QuZhang
# @date         : 2020-12-4 22:50
# @brief        : 读取猫狗数据集
"""
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
from torchvision.transforms import transforms
from tools.my_dataset import CatDogDataset
from torch.utils.data import DataLoader
from tools.common_tools import transform_invert
import matplotlib.pylab as plt


BATCH_SIZE = 5

if __name__ == '__main__':
    # 1.设置数据所在路径
    cats_dogs_dir = os.path.join(BASE_DIR, '..', '..', '..', 'data', 'cats_dogs',)
    train_dir = os.path.join(cats_dogs_dir, 'train')
    test_dir = os.path.join(cats_dogs_dir, 'test1')

    # 2.设置预处理方法
    norm_mean = [0.485, 0.456, 0.406]
    norm_std = [0.229, 0.224, 0.225]
    train_transform = transforms.Compose([
        transforms.Resize((32, 32)),
        transforms.RandomCrop(32, padding=4),
        transforms.ToTensor(),

    ])

    # 3.构建数据集对象
    train_data = CatDogDataset(data_dir=train_dir, transform=train_transform)

    # 4.构建加载数据集的对象
    train_loader = DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)

    # 5. 读取数据集
    for i, data in enumerate(train_loader):
        if i < 1:
            inputs, labels = data
            # print("inputs: ", inputs)
            print("labels: ", labels)
            for j in range(BATCH_SIZE):
                # 将张量转为图像，显示
                img_tensor = inputs[j, ...]
                img = transform_invert(img_tensor, train_transform)
                plt.imshow(img)
                plt.show()
                plt.pause(0.5)
                plt.close()